import cv2
import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.distance import cdist
import matplotlib.patches as patches


def brightest_symmetric_spot(img,
                             inner_radius=15,
                             outer_frac=0.4,
                             thres_rel=0.2,
                             pixel_tol=2,
                             show_plots=True):
    """
    直接对传入的二维灰度图 numpy 数组做 FFT 滤波，仅保留最亮的一对
    中心对称亮点，返回 (亮点坐标对, 重建条纹图)。
    如果失败返回 (None, None)。

    改进：在频域亮点图上标出选取的亮点
    """
    # ------- 1. 输入检查 -------
    if img.ndim != 2:
        raise ValueError('输入图像必须是 2-D 灰度图')
    roi_img = img.astype(np.float32)

    h, w = roi_img.shape
    if h < 2 * inner_radius or w < 2 * inner_radius:
        print('图像太小，无法满足 inner_radius 条件')
        return None, None

    # ------- 2. FFT + 中心化 -------
    f = np.fft.fft2(roi_img)
    fshift = np.fft.fftshift(f)
    mag = np.abs(fshift)
    log_mag = np.log1p(mag)  # 对数变换增强可视化

    # ------- 3. 环形掩码 -------
    cy, cx = h // 2, w // 2
    Y, X = np.ogrid[:h, :w]
    dist = np.sqrt((X - cx) ** 2 + (Y - cy) ** 2)
    outer_radius = min(cx, cy) * outer_frac
    mask = (dist > inner_radius) & (dist < outer_radius)
    mag_band = mag * mask

    # ------- 4. 二值化 -------
    mag_norm = (mag_band - mag_band.min()) / (mag_band.max() - mag_band.min() + 1e-8)
    _, bw = cv2.threshold(mag_norm.astype(np.float32),
                          thres_rel, 1, cv2.THRESH_BINARY)

    # ------- 5. 连通域 → 亮点坐标 -------
    bw_u8 = (bw * 255).astype(np.uint8)
    n, labels, stats, centroids = cv2.connectedComponentsWithStats(bw_u8)
    areas = stats[:, cv2.CC_STAT_AREA]
    idx = np.argsort(areas)[::-1][1:]  # 去掉背景
    pts = centroids[idx]

    # 5.5 去掉离中心太近的亮点（含原点）
    dists_to_center = cdist(pts, [[cx, cy]])
    pts = pts[dists_to_center[:, 0] > inner_radius]
    if len(pts) < 2:
        print('过滤中心后未找到足够亮点！')
        return None, None

    # ------- 6. 对称配对 -------
    center = np.array([cx, cy])
    pts_centered = pts - center
    pairs = []
    used = set()
    for i, p in enumerate(pts_centered):
        if i in used:
            continue
        p_sym = -p
        dists = cdist([p_sym], pts_centered)[0]
        j = np.argmin(dists)
        if dists[j] < pixel_tol and j not in used:
            pairs.append((pts[i], pts[j]))
            used.update([i, j])

    # ===== 新增：打印所有对称亮点对 =====
    print(f"找到 {len(pairs)} 组对称亮点对：")
    for i, (p1, p2) in enumerate(pairs):
        print(f"  对称对 #{i + 1}:")
        print(f"    → 点1: ({p1[0]:.1f}, {p1[1]:.1f})")
        print(f"    → 点2: ({p2[0]:.1f}, {p2[1]:.1f})")
        print(f"    对称中心: ({cx}, {cy})")
        print(f"    两点间距: {np.linalg.norm(p1 - p2):.1f} 像素")
    # ===============================

    if not pairs:
        print('未检测到对称亮点！')
        return None, None

    # ------- 7. 仅保留最亮的一对 -------
    def pair_intensity(pair):
        p1, p2 = pair
        return mag[int(p1[1]), int(p1[0])] + mag[int(p2[1]), int(p2[0])]

    brightest_pair = max(pairs, key=pair_intensity)
    p1, p2 = brightest_pair

    # ------- 8. 生成频域掩膜 -------
    mask_spots = np.zeros_like(mag, dtype=np.uint8)
    for (xi, yi) in (p1, p2):
        cv2.circle(mask_spots, (int(round(xi)), int(round(yi))), 3, 1, -1)

    f_masked = fshift * mask_spots

    # ------- 9. 反变换 -------
    img_back = np.real(np.fft.ifft2(np.fft.ifftshift(f_masked)))

    # ------- 10. 改进的可视化：在频域图上标出亮点 -------
    if show_plots:
        plt.figure(figsize=(15, 5))

        # 原始图像
        plt.subplot(131)
        plt.title('Original Image')
        plt.imshow(roi_img, cmap='gray')
        plt.axis('off')

        # 频域图（对数变换）
        plt.subplot(132)
        plt.title('FFT Magnitude (log scale)')
        plt.imshow(log_mag, cmap='viridis')
        plt.colorbar()

        # 在频域图上标出中心点和环形区域
        plt.scatter(cx, cy, s=50, c='red', marker='+', label='Center')
        inner_circle = plt.Circle((cx, cy), inner_radius, color='red', fill=False, linestyle='--', label='Inner Radius')
        outer_circle = plt.Circle((cx, cy), outer_radius, color='blue', fill=False, linestyle='--',
                                  label='Outer Radius')
        plt.gca().add_patch(inner_circle)
        plt.gca().add_patch(outer_circle)

        # 标出所有检测到的点
        plt.scatter(pts[:, 0], pts[:, 1], s=40, c='cyan', marker='o', alpha=0.7, label='Detected Points')

        # 标出最终选择的对称点对
        plt.scatter([p1[0], p2[0]], [p1[1], p2[1]], s=100, c='red', marker='*', edgecolor='yellow', linewidth=1,
                    label='Selected Pair')

        # 添加对称线
        plt.plot([p1[0], p2[0]], [p1[1], p2[1]], 'r--', linewidth=1, alpha=0.7)

        plt.legend(loc='upper right', fontsize=8)
        plt.axis('off')

        # 重建的条纹图
        plt.subplot(133)
        plt.title('Reconstructed Stripes')
        plt.imshow(img_back, cmap='gray')
        plt.axis('off')

        plt.tight_layout()
        plt.show()

    return brightest_pair, img_back
